# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/11 11:20 
@Author  : ZhangShenao 
@File    : 3.使用Runnable生命周期监听器.py 
@Desc    : 使用Runnable生命周期监听器

在LangChain中,除了在链执行时传递config+callbacks来对链进行监听之外
Runnable还提供了一个简单的方法with_listeners来设置监听器
可以监听开始、结束和异常这3个常见的生命周期
并且with_listeners提供的方法比CallbackHandler更简洁、更统一

with_listeners底层本质上还是只用的CallbackHandler
"""
import time

from langchain_core.runnables import RunnableLambda, RunnableConfig
from langchain_core.tracers import Run


def accumulate(n: int) -> int:
    """业务执行函数"""

    print('-------------accumulate--------------')
    time.sleep(2)
    res = 0
    for i in range(0, n + 1):
        res += i
    return res


def on_start(run_obj: Run, config: RunnableConfig) -> None:
    """
    开始执行监听回调
    :param run_obj: 运行时对象,包含了全部运行时相关信息
    :param config: 运行配置,即调用Runnable组件时传入的配置信息。参数名必须使用config
    """

    print('-------------on_start--------------')
    print(f'run_obj: {run_obj}')
    print(f'config: {config}')


def on_end(run_obj: Run, config: RunnableConfig) -> None:
    """
    结束执行监听回调
    :param run_obj: 运行时对象,包含了全部运行时相关信息
    :param config: 运行配置,即调用Runnable组件时传入的配置信息
    """

    print('-------------on_end--------------')
    print(f'run_obj: {run_obj}')
    print(f'config: {config}')


def on_error(run_obj: Run, config: RunnableConfig) -> None:
    """
    执行异常监听回调
    :param run_obj: 运行时对象,包含了全部运行时相关信息
    :param config: 运行配置,即调用Runnable组件时传入的配置信息
    """

    print('-------------on_error--------------')
    print(f'run_obj: {run_obj}')
    print(f'run_config: {config}')


# 创建Runnable,并设置生命周期监听器
chain = RunnableLambda(accumulate).with_listeners(
    on_start=on_start,
    on_end=on_end,
    on_error=on_error,
)

# 调用chain,并传入运行时配置
result = chain.invoke(10, config={
    'configurable': {
        'executed_by': 'zsa',
    }
})

print(result)
